from typing import Union, Dict

import numpy as np

class Jitter:
    """
    Add gaussian noise to the sample.
    """
    def __init__(self, sigma: float = 0.2, amplitude: float = 0.6, keys=None) -> None:
        self.sigma = sigma
        self.amplitude = amplitude
        self.keys = keys

    def _apply(self, sample: np.ndarray) -> np.ndarray:
        amplitude = self.amplitude * sample
        noise = np.random.normal(loc=0.0, scale=self.sigma, size=sample.shape).astype(np.float32)
        return sample + amplitude * noise

    def __call__(self, data: Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]) -> Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]:
        if isinstance(data, dict):
            for key in self.keys:
                sample = data[key]
                data[key] = self._apply(sample)
        elif isinstance(data, np.ndarray):
            data = self._apply(data)
        return data
